from keras.datasets import mnist
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense,  Conv2D, MaxPooling2D, Flatten

(X_train, y_train), (X_test, y_test) = mnist.load_data()

# 数据预处理
X_train = X_train.reshape(-1, 28, 28, 1)  # normalize
X_test = X_test.reshape(-1, 28, 28, 1)      # normalize
# 讲述像素值转化到0-1区间
X_train = X_train / 255
X_test = X_test / 255
# 将类别向量转换为二进制（只有0和1）的矩阵类型表示。下面是将其分为10个类别
y_train = np_utils.to_categorical(y_train, num_classes=10)
y_test = np_utils.to_categorical(y_test, num_classes=10)


# 下面是模型结构
model = Sequential()
# 卷积层，5x5卷积核，一共20个，激活函数使用relu
model.add(Conv2D(input_shape=(28, 28, 1), kernel_size=(5, 5), filters=20, activation='relu'))
# 最大池化，2x2核，步伐为2，池化后和原来的大小一样（same），所以需要进行填充
model.add(MaxPooling2D(pool_size=(2,2), strides=2, padding='same'))

model.add(Conv2D(kernel_size=(5, 5), filters=50,  activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2,2), strides=2, padding='same'))

# 将上面最后一个输出展开成一个以为数组
model.add(Flatten())
# 全连接层
model.add(Dense(500, activation='relu'))
model.add(Dense(10, activation='softmax'))
# 优化器使用：rmsprop，损失函数使用多分类常用的categorical_crossentropy：，并且记录准确率
model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])

print('Training')
# 训练，整个数据集用来训练两次，每次输入32个数据
model.fit(X_train, y_train, epochs=2, batch_size=32)

print('\nTesting')
loss, accuracy = model.evaluate(X_test, y_test)

print('\ntest loss: ', loss)
print('\ntest accuracy: ', accuracy)